#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import cv2
import numpy as np
import matplotlib.pyplot as plt
import random
import sys
from pathlib import Path
import argparse


class AnnotationVisualizer:
    """
    AnnotationVisualizer
    ===================

    功能：
        可视化YOLO格式的标注文件，显示原始图像和彩色掩码。

    输入：
        - project_path: 项目路径（包含images/、labels/、notes.json）

    输出：
        - 交互式可视化界面

    使用方法：
        visualizer = AnnotationVisualizer(project_path)
        visualizer()
    """

    def __init__(self, project_path, seed=42):
        self.project_path = Path(project_path)
        self.seed = seed

        random.seed(seed)
        np.random.seed(seed)

        self.images_dir = self.project_path / "images"
        self.labels_dir = self.project_path / "labels"
        self.notes_path = self.project_path / "notes.json"

        # 读取标签映射
        self.id_to_name, self.name_to_color = self._load_label_mappings()

        # 获取所有图片文件
        self.all_image_files = sorted([
            f for f in os.listdir(self.images_dir)
            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
        ])

    def _load_label_mappings(self):
        """从notes.json中读取标签映射和颜色映射"""
        if not self.notes_path.exists():
            raise FileNotFoundError(f"未找到 {self.notes_path}")

        with open(self.notes_path, 'r', encoding='utf-8') as f:
            notes_data = json.load(f)

        # 创建ID到名称的映射
        id_to_name = {
            category["id"]: category["name"]
            for category in notes_data["categories"]
        }

        # 定义颜色映射
        name_to_color = {
            "Ripe": (255, 0, 0),
            "Ripe7": (236, 85, 53),
            "Ripe4": (255, 77, 0),
            "Ripe2": (22, 152, 43),
            "Unripe": (0, 100, 0),
            "Disease": (153, 153, 51),
            "strawberry_flower": (255, 215, 0),
            "Stem": (10, 245, 10)
        }

        return id_to_name, name_to_color

    def _generate_random_color(self):
        """生成随机颜色"""
        return (
            random.randint(50, 255),
            random.randint(50, 255),
            random.randint(50, 255)
        )

    def _wait_for_key(self):
        """等待用户按键，返回按下的键"""
        print("按 'n' 或空格继续下一张，按 'q' 退出...")

        # 创建图形窗口并连接按键事件
        fig = plt.gcf()
        key_pressed = None

        def on_key(event):
            nonlocal key_pressed
            key_pressed = event.key
            plt.close(fig)

        # 连接按键事件
        cid = fig.canvas.mpl_connect('key_press_event', on_key)

        # 显示图形并等待
        plt.show(block=True)

        # 断开事件连接
        fig.canvas.mpl_disconnect(cid)

        return key_pressed

    def _process_image_annotations(self, image_file):
        """处理单张图像的标注"""
        # 读取图像
        img_path = self.images_dir / image_file
        img = cv2.imread(str(img_path))

        if img is None:
            print(f"错误: 无法读取图像 {image_file}")
            return None, None, None

        # 转换颜色空间
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_h, img_w = img.shape[:2]

        # 创建彩色掩码
        color_mask = np.zeros((img_h, img_w, 3), dtype=np.uint8)

        # 读取对应的标注文件
        label_file = img_path.stem + ".txt"
        label_path = self.labels_dir / label_file

        polygons = []

        if label_path.exists():
            with open(label_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 解析每行标注
            for line in lines:
                parts = line.strip().split()
                if len(parts) < 3:  # 至少需要类别ID和一个点
                    continue

                class_id = int(parts[0])
                points = [float(x) for x in parts[1:]]

                # 将归一化坐标转换为像素坐标
                pixel_points = []
                for i in range(0, len(points), 2):
                    if i + 1 < len(points):
                        x = int(points[i] * img_w)
                        y = int(points[i + 1] * img_h)
                        pixel_points.append((x, y))

                if len(pixel_points) >= 3:  # 多边形至少需要3个点
                    polygons.append({
                        'class_id': class_id,
                        'points': pixel_points,
                        'color': self._generate_random_color()  # 为每个多边形实例生成随机颜色
                    })

        # 绘制所有多边形
        for polygon in polygons:
            class_id = polygon['class_id']
            points = polygon['points']
            color = polygon['color']

            # 绘制多边形到原图
            pts = np.array(points, np.int32)
            cv2.polylines(img, [pts], True, color, 2)

            # 在彩色掩码上填充多边形
            cv2.fillPoly(color_mask, [pts], color)

            # 添加标签文本
            if points:
                cv2.putText(
                    img,
                    self.id_to_name[class_id],
                    (points[0][0], points[0][1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    color,
                    2
                )

        return img, color_mask, len(polygons)

    def visualize_annotations(self):
        """可视化所有标注"""
        print(f"开始可视化标注 ...")
        print(f"项目路径: {self.project_path}")
        print(f"总共 {len(self.all_image_files)} 张图像")

        for idx, image_file in enumerate(self.all_image_files):
            # 处理图像标注
            img, color_mask, polygon_count = self._process_image_annotations(image_file)

            if img is None:
                continue

            # 显示原始图像和彩色掩码
            plt.figure(figsize=(12, 6))
            plt.subplot(1, 2, 1)
            plt.imshow(img)
            plt.title(f"Image {idx + 1}/{len(self.all_image_files)}: {image_file}")
            plt.axis('off')

            plt.subplot(1, 2, 2)
            plt.imshow(color_mask)
            plt.title(f"Color Mask {idx + 1} ({polygon_count} polygons)")
            plt.axis('off')

            plt.tight_layout()

            # 显示图像信息
            print(f"图像 {idx + 1}/{len(self.all_image_files)}: {image_file} ({polygon_count} 个多边形)")

            # 等待用户按键
            key = self._wait_for_key()

            # 处理按键
            if key in ['q', 'Q', 'escape']:
                print("用户退出程序")
                sys.exit(0)
            elif key in ['n', 'N', ' ', 'enter']:
                # 继续下一张
                plt.close('all')
                continue

        print("所有图像处理完成。")

    def __call__(self):
        self.visualize_annotations()


def main():
    parser = argparse.ArgumentParser(description='YOLO标注可视化工具')
    parser.add_argument('project_path', type=str, help='项目路径（包含images/、labels/、notes.json）')

    args = parser.parse_args()

    visualizer = AnnotationVisualizer(args.project_path)
    visualizer()


if __name__ == "__main__":
    main()